import time
import os
import numpy as np
import cv2
from rknnlite.api import RKNNLite

OBJ_THRESH = 0.25  # 置信度阀值
NMS_THRESH = 0.45  # NMS非极大值抑制阈值
IMG_SIZE = 640

CLASSES = ("person", "bicycle", "car", "motorbike ", "aeroplane ", "bus ", "train", "truck ", "boat", "traffic light",
           "fire hydrant", "stop sign ", "parking meter", "bench", "bird", "cat", "dog ", "horse ", "sheep", "cow", "elephant",
           "bear", "zebra ", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite",
           "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup", "fork", "knife ",
           "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza ", "donut", "cake", "chair", "sofa",
           "pottedplant", "bed", "diningtable", "toilet ", "tvmonitor", "laptop	", "mouse	", "remote ", "keyboard ", "cell phone", "microwave ",
           "oven ", "toaster", "sink", "refrigerator ", "book", "clock", "vase", "scissors ", "teddy bear ", "hair drier", "toothbrush ")

def nms_boxes(boxes, scores):
    """
    非极大值抑制（NMS）：去除重叠度高的冗余检测框，保留置信度最高的框
    Args:
        boxes: 检测框坐标数组，形状为[N, 4]，格式为[x1, y1, x2, y2]
               （x1,y1：左上角坐标；x2,y2：右下角坐标）
        scores: 每个检测框的置信度数组，形状为[N,]
    Returns:
        keep: 保留的检测框索引数组，形状为[K,]（K为筛选后保留的框数量）
    """
    # 提取所有框的左上角x、y坐标
    x = boxes[:, 0]
    y = boxes[:, 1]
    # 计算每个框的宽（x2 - x1）和高（y2 - y1）
    w = boxes[:, 2] - boxes[:, 0]
    h = boxes[:, 3] - boxes[:, 1]

    # 计算每个框的面积（宽×高）
    areas = w * h
    # 对置信度进行降序排序，获取排序后的索引（从高到低）
    order = scores.argsort()[::-1]

    # 保存最终保留的框的索引
    keep = []
    # 循环处理所有框，直到没有剩余框需要处理
    while order.size > 0:
        # 取当前置信度最高的框的索引
        i = order[0]
        keep.append(i)  # 将该框加入保留列表

        # 计算当前最高置信度框与其他所有框的重叠区域
        # 重叠区域左上角x坐标：取两个框左上角x的最大值
        xx1 = np.maximum(x[i], x[order[1:]])
        # 重叠区域左上角y坐标：取两个框左上角y的最大值
        yy1 = np.maximum(y[i], y[order[1:]])
        # 重叠区域右下角x坐标：取两个框右下角x的最小值
        xx2 = np.minimum(x[i] + w[i], x[order[1:]] + w[order[1:]])
        # 重叠区域右下角y坐标：取两个框右下角y的最小值
        yy2 = np.minimum(y[i] + h[i], y[order[1:]] + h[order[1:]])

        # 计算重叠区域的宽和高（若为负则取0，代表无重叠）
        w1 = np.maximum(0.0, xx2 - xx1 + 0.00001)  # +0.00001避免除以0
        h1 = np.maximum(0.0, yy2 - yy1 + 0.00001)
        # 计算重叠区域面积
        inter = w1 * h1

        # 计算交并比（IoU）：重叠面积 /（当前框面积 + 其他框面积 - 重叠面积）
        ovr = inter / (areas[i] + areas[order[1:]] - inter)
        # 保留交并比小于等于NMS阈值的框（重叠度低的框）
        inds = np.where(ovr <= NMS_THRESH)[0]
        # 更新排序索引，只保留符合条件的框（+1是因为排除了当前最高置信度框）
        order = order[inds + 1]

    # 将保留的索引转换为numpy数组并返回
    keep = np.array(keep)
    return keep


def draw(image, boxes, scores, classes):
    """
    在图像上绘制检测结果：包括边界框和类别+置信度标签
    Args:
        image: 原始图像（OpenCV格式，BGR通道），将在该图像上绘制
        boxes: 检测框坐标数组，形状为[K, 4]，格式为[x1, y1, x2, y2]
        scores: 每个检测框的置信度数组，形状为[K,]
        classes: 每个检测框的类别索引数组，形状为[K,]
    """
    # 遍历每个检测框、对应的置信度和类别
    for box, score, cl in zip(boxes, scores, classes):
        # 解析框坐标（top=x1, left=y1, right=x2, bottom=y2）
        top, left, right, bottom = box
        # 打印检测信息（类别名称和置信度）
        print('class: {}, score: {}'.format(CLASSES[cl], score))
        # 打印框坐标信息
        print('box coordinate left,top,right,down: [{}, {}, {}, {}]'.format(top, left, right, bottom))
        
        # 将坐标转换为整数（OpenCV绘图函数要求整数坐标）
        top = int(top)
        left = int(left)
        right = int(right)
        bottom = int(bottom)

        # 绘制边界框：参数依次为图像、左上角坐标、右下角坐标、颜色（BGR）、线宽
        cv2.rectangle(image, (top, left), (right, bottom), (255, 0, 0), 2)
        # 绘制类别+置信度标签：参数依次为图像、文本内容、文本位置、字体、字体大小、颜色、线宽
        cv2.putText(image, '{0} {1:.2f}'.format(CLASSES[cl], score),
                    (top, left - 6),  # 文本位置：框上方6像素处
                    cv2.FONT_HERSHEY_SIMPLEX,  # 字体（仅支持英文）
                    0.6, (0, 0, 255), 2)  # 字体大小0.6，红色（BGR），线宽2


def boxes_computer(image, xywh):
    """
    计算检测框在原始图像上的坐标（将模型输出的归一化坐标转换为实际像素坐标）
    Args:
        image: 原始图像（用于获取图像宽高）
    Returns:
        boxes: 转换后的检测框坐标数组，形状为[K, 4]，格式为[x1, y1, x2, y2]
        classes: 筛选后的类别索引数组，形状为[K,]
        scores: 筛选后的置信度数组，形状为[K,]
    """
    # 获取原始图像的高和宽（用于坐标转换）
    img_h_orginal, img_w_orginal = image.shape[:2]
    # 模型输入图像的尺寸（固定为IMG_SIZE×IMG_SIZE）
    img_w_infer, img_h_infer = IMG_SIZE, IMG_SIZE

    # -------------------------- 坐标转换核心逻辑 --------------------------
    # 1. 将模型输出的原始xywh（基于640×640输入图像）转换到原始图像尺寸
    # cx, cy：目标中心坐标；w, h：目标宽高
    # 转换公式：原始图像坐标 =（模型输出坐标 / 模型输入尺寸）× 原始图像尺寸
    cx = (xywh[:, 0] / img_w_infer) * img_w_orginal  # 中心x坐标转换
    cy = (xywh[:, 1] / img_h_infer) * img_h_orginal  # 中心y坐标转换
    w = (xywh[:, 2] / img_w_infer) * img_w_orginal   # 宽度转换
    h = (xywh[:, 3] / img_h_infer) * img_h_orginal   # 高度转换

    # 2. 将中心坐标+宽高（xywh）转换为左上角+右下角（x1y1x2y2）
    x1 = cx - w/2  # 左上角x = 中心x - 宽/2
    y1 = cy - h/2  # 左上角y = 中心y - 高/2
    x2 = cx + w/2  # 右下角x = 中心x + 宽/2
    y2 = cy + h/2  # 右下角y = 中心y + 高/2
    
    # 3. 堆叠坐标为[N, 4]的检测框数组
    boxes = np.stack([x1, y1, x2, y2], axis=1)  # 形状: (N, 4)

    # 4. 应用非极大值抑制（NMS）去除重叠框
    keep = nms_boxes(boxes, conf_scores)
    # 5. 根据NMS结果筛选框、类别和置信度
    boxes = boxes[keep]       # 筛选后的检测框

    return keep, boxes
  
if __name__ == '__main__':
    """
    主程序入口：
    流程：1. 初始化RKNN对象 → 2. 加载模型 → 3. 初始化NPU运行时 → 4. 图像预处理 → 
         5. 模型推理 → 6. 输出后处理（坐标解析+筛选） → 7. 结果绘制与保存 → 8. 资源释放
    """
    # -------------------------- 1. 路径初始化（获取文件绝对路径） --------------------------
    # 获取当前脚本所在目录的绝对路径（避免相对路径报错）
    pwd_dir = os.path.dirname(os.path.abspath(__file__))
    # 拼接路径
    image_path = os.path.join(pwd_dir, 'bus.jpg')  # 原始图片路径
    rknn_model = os.path.join(pwd_dir, 'yolov8n_rk3576.rknn')  # 模型路径

    # -------------------------- 2. 计时初始化（统计各环节耗时） --------------------------
    start_time = time.time()  # 总计时起点

    # -------------------------- 3. 创建RKNN对象（初始化模型载体） --------------------------
    rknn_lite = RKNNLite()  # 创建RKNN Lite轻量版对象（适配嵌入式设备）
    create_time = time.time()  # 记录对象创建完成时间
    # 打印“对象创建”耗时（毫秒级，保留2位小数）
    print(f"创建耗时: {(create_time - start_time)*1000:.2f} ms")

    # -------------------------- 4. 加载RKNN模型（读取模型权重与结构） --------------------------
    print(f"当前加载的模型路径: {rknn_model}")  # 打印模型路径（调试用，确认路径正确）
    # 加载RKNN模型，返回状态码（0表示成功，非0表示失败）
    ret = rknn_lite.load_rknn(rknn_model)
    if ret != 0:
        print('Load RKNN model failed（模型加载失败）')
        exit(ret) 

    load_time = time.time()  # 记录模型加载完成时间
    print(f"加载耗时: {(load_time - create_time)*1000:.2f} ms")

    # -------------------------- 5. 初始化NPU运行时（激活硬件计算资源） --------------------------
    # 初始化RKNN运行时环境，指定使用NPU核心0
    ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
    run_time = time.time()  # 记录运行时初始化完成时间
    # 打印“运行时初始化”耗时（此环节耗时较长，与硬件资源分配相关）
    print(f"初始化耗时: {(run_time - load_time)*1000:.2f} ms")

    # 检查运行时初始化是否成功，失败则退出
    if ret != 0:
        print('Init runtime environment failed!（运行时环境初始化失败）')
        exit(ret)

    # -------------------------- 6. 图像预处理（适配模型输入要求） --------------------------
    # 6.1 读取原始图像（OpenCV默认BGR通道）
    img_original = cv2.imread(image_path)  # 读取原始图像（用于后续原始尺寸坐标转换）
    # 6.2 BGR转RGB（YOLOv8训练/推理使用RGB通道，需与模型输入格式对齐）
    img_original = cv2.cvtColor(img_original, cv2.COLOR_BGR2RGB)  # BGR→RGB通道转换

    # 6.3 准备“绘图用图像”（3维uint8格式，不添加Batch维度，用于最终结果绘制）
    # 缩放原始图像至模型输入尺寸（640x640），保持uint8类型（OpenCV绘图要求）
    img_draw = cv2.resize(img_original, (IMG_SIZE, IMG_SIZE))  # 输出形状：(640, 640, 3)

    # 6.4 准备“推理用图像”（4维float32格式，满足模型输入要求）
    img_infer = img_draw.copy()  # 基于绘图图像复制（避免修改原始绘图数据）
    # 添加Batch维度（模型要求输入格式为[Batch, H, W, C]，单图推理Batch=1）
    img_infer = np.expand_dims(img_infer, axis=0).astype(np.float32)  # 输出形状：(1, 640, 640, 3)
    # 归一化：将像素值从[0,255]转为[0,1]（YOLOv8训练时的预处理逻辑，推理必须对齐）
    img_infer = img_infer / 255.0  # 关键步骤：不归一化会导致推理结果异常

    image_time = time.time()  # 记录图像预处理完成时间
    print(f"图形处理耗时: {(image_time - run_time)*1000:.2f} ms")

    # -------------------------- 7. 模型推理（NPU执行目标检测计算） --------------------------
    # 执行推理：inputs参数需为列表格式（即使单图也需用[]包裹，符合RKNN API要求）
    outputs = rknn_lite.inference(inputs=[img_infer])
    inference_time = time.time()  # 记录推理完成时间
    print(f"推理耗时: {(inference_time - image_time)*1000:.2f} ms")

    # -------------------------- 8. 推理输出后处理（解析坐标与筛选有效框） --------------------------
    # 8.1 提取模型输出张量（YOLOv8仅1个输出张量，形状为[1, 84, 8400]）
    output = outputs[0]  # 取出列表中唯一的输出张量

    # 8.2 简化输出格式（去除无用维度+转置，便于后续处理）
    # np.squeeze：删除Batch维度（从[1,84,8400]→[84,8400]）
    # .T：转置为[8400, 84]（每行对应1个锚框，每列对应1个预测参数）
    predictions = np.squeeze(output).T  # 最终形状: (8400, 84)（8400个锚框，每个84个参数）

    # 8.3 拆分预测参数（坐标+类别得分）
    xywh = predictions[:, :4]  # 前4列：每个锚框的原始坐标（x中心, y中心, 宽, 高）→ 未激活
    class_scores = predictions[:, 4:]  # 后80列：每个锚框在80个类别上的得分（COCO数据集）

    # 8.4 计算置信度与类别ID
    conf_scores = class_scores.max(axis=1)  # 每个锚框的“最大类别得分”→ 即该框的置信度
    class_ids = class_scores.argmax(axis=1)  # 每个锚框“得分最高的类别索引”→ 即目标类别

    # 8.5 过滤低置信度框（去除置信度<阈值的无效框，减少后续计算量）
    mask = conf_scores > OBJ_THRESH  # 生成布尔掩码：True表示置信度达标（>0.25）
    # 根据掩码筛选有效数据
    xywh = xywh[mask]          # 筛选后的坐标
    conf_scores = conf_scores[mask]  # 筛选后的置信度
    class_ids = class_ids[mask]      # 筛选后的类别ID

    # 8.6 检查是否有有效框，无则标记为None；有则计算原始图/缩放图的检测框
    if len(conf_scores) == 0:
        print("未检测到目标（所有框的置信度均低于阈值）")
        boxes, classes, scores = None, None, None
    else:
        # 调用boxes_computer函数：将锚框坐标转换为图像像素坐标（适配原始图和缩放图）
        keep_ori, boxes_ori = boxes_computer(img_original, xywh)  # 原始尺寸图像的检测框
        classes_ori = class_ids[keep_ori] # 筛选后的类别索引
        scores_ori = conf_scores[keep_ori]# 筛选后的置信度
        kepp_draw, boxes_draw = boxes_computer(img_draw, xywh)   # 640x640缩放图的检测框
        classes_draw = class_ids[kepp_draw] # 筛选后的类别索引
        scores_draw = conf_scores[kepp_draw]# 筛选后的置信度

    # -------------------------- 9. 结果绘制与保存（可视化检测结果） --------------------------
    # 9.1 通道转换：将RGB格式转为BGR（OpenCV保存/显示图像默认BGR通道，不转换会导致颜色异常）
    img_ori = cv2.cvtColor(img_original, cv2.COLOR_RGB2BGR)  # 原始图：RGB→BGR
    img_draw = cv2.cvtColor(img_draw, cv2.COLOR_RGB2BGR)     # 缩放图：RGB→BGR

    # 9.2 绘制检测结果（仅当有有效框时绘制）
    if boxes_ori is not None:
        # 在原始尺寸图像上绘制框+标签
        draw(img_ori, boxes_ori, scores_ori, classes_ori) 
    if boxes_draw is not None:
        # 在640x640缩放图像上绘制框+标签
        draw(img_draw, boxes_draw, scores_draw, classes_draw)  

    # 9.3 保存绘制结果到本地（路径为当前脚本目录，文件名分别为out_ori.jpg/out_draw.jpg）
    cv2.imwrite((pwd_dir + "/out_ori.jpg"), img_ori)  # 保存原始尺寸结果图
    cv2.imwrite((pwd_dir + "/out_draw.jpg"), img_draw)# 保存缩放尺寸结果图
    print(f"结果保存目录: {pwd_dir}（可在该目录查看out_ori.jpg和out_draw.jpg）")

    # -------------------------- 10. 计时与资源释放（统计总耗时+释放NPU资源） --------------------------
    top_time = time.time()  # 记录后处理（绘制+保存）完成时间
    print(f"后处理耗时: {(top_time - inference_time)*1000:.2f} ms")
    print(f"总耗时: {(top_time - create_time)*1000:.2f} ms")

    labels = []
    for box, score, cl in zip(boxes_ori, scores_ori, classes_ori):
        top, left, right, bottom = box

        labels.append({
            "label": f"{CLASSES[cl]}",
            "category": "物体",
            "confidence": float(score),
            "bbox": [float(top), float(left), float(right), float(bottom)],
            "description": f"检测到{CLASSES[cl]}"
        })
    print(labels)

    # 释放RKNN资源（关闭NPU连接，避免占用硬件资源）
    rknn_lite.release()
